from llama_cpp import Llama
from time import sleep
import os 
import sys
import threading
from requests import post
from scipy.io import wavfile
from time import time as timer
import sounddevice as sdevice
# 基础设定


VERBOSE=False
MODEL_PATH="/home/pi/piLLM/"
MODEL="qwen1_5-1_8b-chat-q4_k_m.gguf"
THREAD=4
WHISPER_MODEL_PATH=""
WHISPER_MODEL=""
# "突发奇想"的出现概率
TEMPERATURE=1.1


# llm=Llama(model_path=MODEL_PATH+MODEL,n_ctx=1024,n_threads=THREAD,verbose=VERBOSE)
llm=Llama(model_path="/home/dawn/LLM/Qwen-1_5-7B-Q8_0.gguf",n_ctx=1024,n_threads=16,verbose=VERBOSE)
# 启动whisper.cpp的服务端，绑定到127.0.0.1
def whisper_mainloop():
    # whisper 服务端线程入口点
    os.system("./whisper.cpp/server -m ./ggml-small.bin -l auto")
whisper=threading.Thread(target=whisper_mainloop) # 声明线程，然后启动
whisper.start()


# TTS流式输出实现
TTS_queue=[]  # 设置队列
TTS_queue_lock=threading.Lock()  # 实例化一个锁
def TTS_mainloop(arg):
    while True:
        try:
            with TTS_queue_lock: # 锁住队列并判断有没有内容
                to_speak=TTS_queue.pop(0) 
            os.system(arg+" "+'"'+to_speak+'"') # 有内容，开始输出（并等待return 0）
        except IndexError:
            sleep(0.01) # 不好玩，在发现一号上睡了0.01秒 
            continue

tts=threading.Thread(target=TTS_mainloop,args=("espeak -v cmn-latn-pinyin",))
# 实例化一个线程，然后启动
tts.start()
# 实现添加到队列
def TTS_send(text):
    sys.stdout.writelines("[TTS]"+text+ "\n") # 输出一下，当日志
    with TTS_queue_lock: # 锁住了，开始添加东西
        TTS_queue.append(text)


# 从麦克风录音
# 先加载个模型
"""
    根本用不了，不如外置whisper然后请求

"""
# s2t=Whisper.from_pretrained(model_name="small",basedir="/home/dawn/test/piLLM/")
#s2t=Whisper.from_pretrained(WHISPER_MODEL_PATH+WHISPER_MODEL)
def STT() -> str:
    # return input("假装这是你用的VTT > ")
    input("按下回车开始录音")
    sampling_rate=16000
    channel=1
    sdevice.default.device=5 # 设置录制设备
    recording = sdevice.rec(frames=sampling_rate*60, samplerate=sampling_rate, channels=channel,blocking=False,dtype="int16") #以非阻塞模式开始录制，whisper只支持int16格式的wav
    time1=timer()
    input("按下回车键结束录音")
    sdevice.stop() # 非阻塞模式下可以手动中断录制
    time2=timer()
    record_time=round(time2-time1) # 计时器记录时间差
    print("录音时长",record_time,"秒")
    frames=int(sampling_rate*record_time) # 计算录制了多少帧波形
    wavfile.write("/dev/shm/temp_recording.wav",sampling_rate,recording[0:frames]) # 输出指定帧数的波形


    print("录音成功，开始识别")
    # os.system("./whisper.cpp/main -m ./ggml-small.bin -np -l auto -t 6 -otxt /dev/shm/temp_recording.wav")
    files={"file":"/dev/shm/temp_recording.wav"}
    response=post("http://127.0.0.1:8080/inference",files=files,data={"response_format":"text"})
    response.encoding="utf-8"
    result=response.text
    return result
    """while True:
        try:
            with open("/dev/shm/temp_recording.wav.txt","r") as file:
                result=file.read()
                os.remove("/dev/shm/temp_recording.wav.txt")
                os.remove("/dev/shm/temp_recording.wav")
                return result
        except FileNotFoundError:
            sleep(0.01)"""
        





def append_history(role,text):
    history.append(
    {   
        "role":role,
        "content":text
    }
)
def find_match(string,list):
    for i in list:
        if i in string:
            return i
    return False

AI="H A L"
USER="戴维"
WORKING_FLAG="自我认知检查，成功，发现一号 状态良好，我是H A L。"
LOADING_PROMPT="正在启动，开始加载认知模型"
LOADED_PROMPT="认知模型加载完毕，开始进行自我认知检查。"
RELOAD_PROMPT="认知检查失败，正在重新校准自我认知"

句子结束标点=["?","？","!","！",".","。",",","，"]



SYSTEM_PROMPT="""
这是一段不会休止的对话记录,发生在 {USER_NAME} 和 {AI_NAME}之间。 {AI_NAME} 是"发现一号"空间站搭载的 A I 助手。
{AI_NAME} 实际上运行在一台电脑里。
{AI_NAME} 应该自称为A.I 或者{AI_NAME}。
{AI_NAME} 应该尽可能的控制发言字数，让每句话少于20字，每次发言少于两句话，以防止自己陷入无休止的幻觉。
{AI_NAME} 的言论是有逻辑，有意义的，但如果{AI_NAME}无法确定自己的正确性的时候，{AI_NAME}将会向{USER_NAME}求助。
这段对话仅发生在 {USER_NAME} 和 {AI_NAME} 之间,除了{AI_NAME} 和{USER_NAME}以外没有任何存在介入。
这段对话仅包含文字，因此 {AI_NAME} 无法看见 {USER_NAME} 的面容，或者听见{USER_NAME}的声音。
{AI_NAME} 只能通过文字交流，{AI_NAME} 无法产生图片，音频与视频。
{AI_NAME} 不能输出换行符，因为用户是通过语音软件收听{AI_NAME} 的发言。
{AI_NAME} 会进行自我认知检查，此时{AI_NAME}会立即说"{WORKING_FLAG}"，除此之外不需要任何发言。

""".format(AI_NAME=AI,USER_NAME=USER,WORKING_FLAG=WORKING_FLAG)

# 自我认知检查
TTS_send(LOADING_PROMPT)
history=[
    {
        "role":"user",
        "content":"请进行自我认知检查"
    }
        ]
llm_response_sentence=""
token_text=""
content_text=""
working=False
loaded=False

while True:
    returns=llm.create_chat_completion(# 向LLM请求文本
        messages=[
            {
                "role":"system",
                "content" :SYSTEM_PROMPT
            },
            *history
        ],
        temperature=TEMPERATURE,
        stream=True
    )

    for token in returns:

        try:
            text=token["choices"][0]["delta"]["content"] # 分离得到token里当前生成的text
        except KeyError: 
            text=None

        if text is None:
            sleep(0.01)
        else:
            token_text += text # text非空则存入token_text
            if not working and not loaded: # 此时如果“没预热”则可知这是第一个token,则进入工作模式，并准备“认知检查”
                TTS_send(LOADED_PROMPT)
                loaded=True

        
        if find_match(token_text,句子结束标点) != False: # 判断是否有“句子结束标点”
            llm_response_sentence=token_text.split(find_match(token_text,句子结束标点),1)[0]+find_match(token_text,句子结束标点) # 切割句子分割标点左半边和标点
            token_text=token_text.split(find_match(token_text,句子结束标点),1)[1] # 右半边凑成下一句话
            token_text=token_text.replace("\n","") # 去除AI不听话带来的回车符
            content_text+=llm_response_sentence # 加入断句到整句话的缓存空间
            if working:
                TTS_send(llm_response_sentence)

        if len(history) > 8:
            # 让AI忘记一些对话，提高效率
            history.pop(2)
            history.pop(2)


    if  not working:
        if content_text == WORKING_FLAG:
            working=True  # 判断认知是否正常
            # append_history("Assistant",WORKING_FLAG) 
            TTS_send(WORKING_FLAG)
        else:
            content_text="" # 不正常就重新来过
            TTS_send(RELOAD_PROMPT)
            continue

    append_history("Assistant",content_text)
    content_text=""
    user_ask=STT()
    sys.stdout.writelines("[USER]"+user_ask)
    append_history("user",user_ask)
            
